import psycopg2
import mysql.connector
from psycopg2 import sql

def get_db_connection(config):
    if config.db_type == "postgresql":
        return psycopg2.connect(
            dbname=config.db_name,
            user=config.db_user,
            password=config.db_password,
            host=config.db_host,
            port=config.db_port
        )
    elif config.db_type == "mysql":
        return mysql.connector.connect(
            database=config.db_name,
            user=config.db_user,
            password=config.db_password,
            host=config.db_host,
            port=config.db_port
        )
    else:
        raise ValueError(f"Unsupported database type: {config.db_type}")

def generate_new_table_name(table_name):
    return f"{table_name}_sensitive"

def get_or_create_mapping_id(cursor, new_table_name, column, value):
    # 检查映射表中是否存在相同值，若存在则返回其 ID
    cursor.execute(sql.SQL("""
        SELECT id FROM {} WHERE {} = %s
    """).format(sql.Identifier(new_table_name), sql.Identifier(column)), (value,))
    result = cursor.fetchone()
    if result:
        return result[0]
    # 若不存在，则插入新记录并返回新 ID
    cursor.execute(sql.SQL("""
        INSERT INTO {} ({}) VALUES (%s) RETURNING id
    """).format(sql.Identifier(new_table_name), sql.Identifier(column)), (value,))
    return cursor.fetchone()[0]